import json
import time
import warnings

import pandas as pd
import sqlalchemy
from sklearn.ensemble import AdaBoostRegressor, RandomForestClassifier
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor

from utils.constants import *

warnings.filterwarnings("ignore")
city_map = {}
city_cnt = 0
house_type_map = {}
house_type_cnt = 0
sale_status_map = {}
sale_status_cnt = 0


def preprocess_city(s):
    """
    将城市转换为整数编码
    :param s: 城市
    :return: 整数编码
    """
    global city_cnt, city_map

    if s not in city_map:
        city_map[s] = city_cnt
        city_cnt += 1
    return city_map[s]


def preprocess_house_type(s):
    """
    将房屋类型转换为整数编码
    :param s: 房屋类型
    :return: 整数编码
    """
    global house_type_cnt, house_type_map

    if s not in house_type_map:
        house_type_map[s] = house_type_cnt
        house_type_cnt += 1
    return house_type_map[s]


def preprocess_sale_status(s):
    """
    将销售状态转换为整数编码
    :param s: 房屋销售状态
    :return: 整数编码
    """
    global sale_status_cnt, sale_status_map

    if s not in sale_status_map:
        sale_status_map[s] = sale_status_cnt
        sale_status_cnt += 1
    return sale_status_map[s]


def preprocess_rooms_desc(s):
    """
    将房屋户型描述转换为数值
    :param s: json格式的房屋户型描述
    :return: 数值
    """
    rooms_desc = json.loads(s)
    room = 1
    lazy = 1
    alpha = 0.5
    for r in rooms_desc:
        if r is None or r == "":
            continue
        room = room + int(r) * lazy * alpha
        lazy = lazy * (1 - alpha)
        alpha = alpha * alpha + 0.1
    return room


def preprocess_area_range(s):
    """
    将面积范围转换为数值
    :param s: json格式的面积范围
    :return: 数值
    """
    area_range = json.loads(s)
    area = 0
    if len(area_range) == 1:
        if area_range[0] is None or area_range[0] == "":
            area = 0
        else:
            area = float(area_range[0])
    elif len(area_range) == 2:
        area = (float(area_range[0]) * 1.2 + float(area_range[1]) * 0.8) / 2
    return area


def preprocess_data(data):
    data["city"] = data["city"].apply(preprocess_city)
    data["house_type"] = data["house_type"].apply(preprocess_house_type)
    data["sale_status"] = data["sale_status"].apply(preprocess_sale_status)
    data["rooms_desc"] = data["rooms_desc"].apply(preprocess_rooms_desc)
    data["area_range"] = data["area_range"].apply(preprocess_area_range)


def train_model(model, X, y):
    """
    训练模型
    :param model: 模型
    :param X: 特征
    :param y: 标签
    :return: 训练后的模型
    """
    model.fit(X, y)
    return model


def predict_price(model, city, rooms_desc, area_range, house_type, sale_status):
    """
    预测房屋价格
    :param model:
    :param city:
    :param rooms_desc:
    :param area_range:
    :param house_type:
    :param sale_status:
    :return:
    """
    X_test = pd.DataFrame(
        [[city, rooms_desc, area_range, house_type, sale_status]],
        columns=["city", "rooms_desc", "area_range", "house_type", "sale_status"],
    )
    X_test["city"] = X_test["city"].apply(preprocess_city)
    X_test["house_type"] = X_test["house_type"].apply(preprocess_house_type)
    X_test["sale_status"] = X_test["sale_status"].apply(preprocess_sale_status)
    y_pred = model.predict(X_test)
    return y_pred[0]


# 数据库连接信息
engine_url = (
    f"{DATABASE_TYPE}+{DRIVER}://{USERNAME}:{PASSWORD}@{HOST}:{PORT}/{DATABASE_NAME}"
)
# 连接数据库
engine = sqlalchemy.create_engine(engine_url)
# 读取数据
df = pd.read_sql(
    "SELECT `city`, `rooms_desc`, `area_range`, `house_type`, `sale_status`,`prices` FROM house_info",
    con=engine,
)[["city", "rooms_desc", "area_range", "house_type", "sale_status", "prices"]]

preprocess_data(df)
# print(df)

# 线性回归、AdaBoost、随机森林、决策树
line = LinearRegression()
ada = AdaBoostRegressor(loss="square")
rf = RandomForestClassifier(oob_score=True, min_samples_split=5)
dt = DecisionTreeRegressor()
models = [line, ada, rf, dt]

# 训练模型
X_train, X_test, y_train, y_test = train_test_split(
    df.drop("prices", axis=1), df["prices"], test_size=0.2, random_state=22
)


def evaluate_model(models, X_train, X_test, y_train, y_test):
    """
    评估模型的好坏
    :param models:
    :param X_train:
    :param X_test:
    :param y_train:
    :param y_test:
    :return:
    """
    for model in models:
        start_time = time.time()
        model = train_model(model, X_train, y_train)
        score = model.score(X_test, y_test)
        end_time = time.time()
        print(
            f"{model.__class__.__name__} score: {score}, time: {(round((end_time - start_time) * 1000, 3))}ms"
        )


def getPredict(city, rooms_desc, area_range, house_type, sale_status):
    """
    获取预测房屋价格（可以先选定合适的模型，提高准确率和预测速度）
    :param city:
    :param rooms_desc:
    :param area_range:
    :param house_type:
    :param sale_status:
    :return:
    """
    ans = []
    for model in models:
        model = train_model(model, X_train, y_train)
        score = model.score(X_test, y_test)
        price = predict_price(
            model, city, rooms_desc, area_range, house_type, sale_status
        )
        ans.append([model.__class__.__name__, score, price])
    ans.sort(key=lambda x: x[1], reverse=True)
    return int(ans[0][2])


if __name__ == "__main__":
    evaluate_model(models, X_train, X_test, y_train, y_test)
    print(getPredict("保定", 2, 88, "住宅", "在售"))
